package org.msi.bayesiannetwork;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * <p>Represents a bayesian network.</p>
 *
 * @author Jakub Narloch
 * @version 1.0
 */
public class BayesianNetwork {

    /**
     * <p>Represents the node list.</p>
     */
    private final List<AbstractNode> nodes = new ArrayList<AbstractNode>();

    /**
     * <p>Represents a map of nodes.</p>
     */
    private final Map<String, AbstractNode> nodeMap = new HashMap<String, AbstractNode>();

    /**
     * <p>Represents the edges map, where as keys are used node indexes and value contains list of 'neighbour'
     * nodes.</p>
     */
    private final Map<Integer, List<Integer>> edges = new HashMap<Integer, List<Integer>>();

    /**
     * Retrieves the edge list.
     * @return  Map of edges
     */
    public Map<Integer, List<Integer>> getEdges() {
        return edges;
    }

    /**
     * <p>Adds node with specified name, alias and probability to bayesian network.</p>
     *
     * @param nodeName    the node name
     * @param nodeAlias   the node alias

     * @param probability the node probability
     *
     * @return returns the node index
     */
    public int addNode(String nodeName, String nodeAlias, double probability) {

        return addNode(new Node(nodeName, nodeAlias, probability));
    }

    /**
     * <p>Adds conditional node with specified name, alias and probability to bayesian network.</p>
     *
     * @param nodeName               the node name
     * @param nodeAlias              the node alias
     * @param conditionalProbability the node probability
     *
     * @return the node index
     */
    public int addConditionalNode(String nodeName, String nodeAlias, ConditionalProbability conditionalProbability) {

        return addNode(new ConditionalNode(nodeName, nodeAlias, conditionalProbability));
    }

    /**
     * <p>Adds specified node to bayesian network.</p>
     *
     * @param node the node to add
     *
     * @return the node index
     */
    public int addNode(AbstractNode node) {

        node.setId(nodes.size());
        nodes.add(node);
        nodeMap.put(node.getName(), node);
        return node.getId();
    }

    /**
     * <p>Adds edge connecting two nodes.</p>
     *
     * @param fromNode the starting node
     * @param toNode   the ending node
     */
    public void addEdge(int fromNode, int toNode) {

        if (!edges.containsKey(fromNode)) {
            edges.put(fromNode, new ArrayList<Integer>());
        }

        edges.get(fromNode).add(toNode);
    }

    /**
     * <p>Adds edge connecting two nodes.</p>
     *
     * @param fromNodeName the starting node name
     * @param toNodeName   the ending node name
     *
     * @throws NodeNotFoundException if of the nodes does not exists
     */
    public void addEdge(String fromNodeName, String toNodeName) throws NodeNotFoundException {

        addEdge(getNodeByName(fromNodeName).getId(), getNodeByName(toNodeName).getId());
    }

    /**
     * <p>Adds edge connecting two nodes.</p>
     *
     * @param fromNodeAlias the starting node alias
     * @param toNodeAlias   the ending node alias
     *
     * @throws NodeNotFoundException if one of the nodes does not exists
     */
    public void addEdgeByAlias(String fromNodeAlias, String toNodeAlias) throws NodeNotFoundException {

        addEdge(getNodeByVariableName(fromNodeAlias).getId(), getNodeByVariableName(toNodeAlias).getId());
    }

    /**
     * <p>Finds node by it's name.</p>
     *
     * @param nodeName the node name to find
     *
     * @return the found node, that matches the given name
     *
     * @throws NodeNotFoundException if node with specified name hasn't been found
     */
    public AbstractNode getNodeByName(final String nodeName) throws NodeNotFoundException {

        if (!nodeMap.containsKey(nodeName)) {
            throw new NodeNotFoundException(String.format("The node with name %s does not exists.", nodeName));
        }

        return nodeMap.get(nodeName);
    }

    /**
     * <p>Finds node by it's variable name.</p>
     *
     * @param nodeVariable the node variable name
     *
     * @return the found node, that matches the given variable name
     *
     * @throws NodeNotFoundException if node with specified variable name hasn't been found
     */
    public AbstractNode getNodeByVariableName(final String nodeVariable) throws NodeNotFoundException {

        return findNode(new Criteria<AbstractNode>() {
            public boolean check(AbstractNode obj) {

                return obj.getVariable().equals(nodeVariable);
            }
        });
    }

    /**
     * <p>Finds a node meeting the specified criteria. The criteria are given by instance of passed {@link
     * Criteria}.</p>
     *
     * @param criteria the {@link Criteria} instance
     *
     * @return the found node
     *
     * @throws NodeNotFoundException if node hasn't been found
     */
    public AbstractNode findNode(Criteria<AbstractNode> criteria) throws NodeNotFoundException {

        for (AbstractNode node : nodes) {
            if (criteria.check(node)) {
                return node;
            }
        }

        throw new NodeNotFoundException("The specified node does not exists.");
    }

    /**
     * <p>Retrieves the variable names list.</p>
     *
     * @return the list of variable names
     */
    public List<String> getVariableList() {

        List<String> variables = new ArrayList<String>();

        for (AbstractNode node : nodes) {

            variables.add(node.getVariable());
        }

        return variables;
    }

    /**
     * <p>Retrieves the variable probability.</p>
     *
     * @param variable the variable name
     *
     * @return the variable probability
     *
     * @throws IllegalArgumentException if the observations doesn't contain all the predecessors for the given node
     * @throws IllegalStateException    if the node for the given variable has unknown type
     * @throws NodeNotFoundException    if the variable with the given name doesn't exists
     */
    public double getVariableProbability(Variable variable, List<Variable> observations) throws NodeNotFoundException {

        AbstractNode node = getNodeByVariableName(variable.getVariableName());
        double prop;
        if (node instanceof Node) {

            prop = ((Node) node).getProbability();
            return variable.isPrimitive() ? prop : 1 - prop;

            //return 1D;

        } else if (node instanceof ConditionalNode) {

            ConditionalNode conditionalNode = (ConditionalNode) node;
            ConditionalProbability conditionalProbability = conditionalNode.getConditionalProbability();

            // puts into the known variable all the variable from conditional probability that are present in
            // the observations, any 'unknown' variable is put into unknownVariableNames list
            Map<String, Boolean> variables = new HashMap<String, Boolean>();
            List<String> unknownVariableNames = new ArrayList<String>();
            Variable var;
            for (String varName : conditionalProbability.getVariableNames()) {

                if ((var = findVariable(varName, observations)) != null) {

                    variables.put(var.getVariableName(), var.isPrimitive());
                } else {

                    unknownVariableNames.add(varName);
                }
            }

            if (unknownVariableNames.size() == 0) {
                prop = conditionalProbability.getProbability(variables);
                return variable.isPrimitive() ? prop : 1 - prop;
            } else {

                throw new IllegalArgumentException("Observations doesn't contain all required variables.");
            }
        } else {

            throw new IllegalStateException(String.format("The node of type %s is not supported.",
                    node.getClass().getName()));
        }
    }

    /**
     * <p>Retrieves the list containing all nodes from which the given variable can be reached.</p> <p>Nodes that don't
     * have any predecessors in the graph are excluded.</p>
     *
     * @param variable the variable
     *
     * @return list containing all the nodes
     *
     * @throws NodeNotFoundException if the variable with the given name does not exists
     */
    public List<Variable> getAllVariableParents(Variable variable) throws NodeNotFoundException {

        AbstractNode node;
        List<AbstractNode> nodes = new ArrayList<AbstractNode>();
        List<AbstractNode> parents;

        nodes.add(getNodeByVariableName(variable.getVariableName()));
        List<Variable> variables = new ArrayList<Variable>();
        while (nodes.size() > 0) {
            node = nodes.get(0);
            nodes.remove(0);

            parents = findParents(node);
            nodes.addAll(parents);

            // make sure not to add the variable isn't the starting variable
            if (!node.getVariable().equals(variable.getVariableName())) {

                variables.add(new Variable(node.getVariable(), true));
            }
        }

        // reverse the order of the collections
        Collections.reverse(variables);

        return variables;
    }

    /**
     * <p>Finds the variable with given name in list containing all observations.</p>
     *
     * @param varName      the name of the variable
     * @param observations the list of the observations
     *
     * @return the found variable or null if it hasn't been found
     */
    private Variable findVariable(String varName, List<Variable> observations) {

        for (Variable variable : observations) {
            if (variable.getVariableName().equals(varName)) {
                return variable;
            }
        }

        return null;
    }

    /**
     * <p>Finds all 'parent' nodes for the given node.</p>
     *
     * @param node the node to use
     *
     * @return the list of 'parent' nodes
     */
    private List<AbstractNode> findParents(AbstractNode node) {

        List<AbstractNode> parents = new ArrayList<AbstractNode>();
        for (Map.Entry<Integer, List<Integer>> edge : edges.entrySet()) {
            if (edge.getValue().contains(node.getId())) {
                parents.add(nodes.get(edge.getKey()));
            }
        }

        return parents;
    }

    /**
     * <p>This is template class used for testing if passed object meets the required criteria.</p> <p>It may be used
     * i.e. for searching specified object within a collection.</p>
     *
     * @author Jakub Narloch
     * @version 1.0
     */
    public static interface Criteria<T> {

        /**
         * <p>Checks if passed objects, meets the requirements.</p>
         *
         * @param obj object to check
         *
         * @return true if object meats the requirements, false otherwise
         */
        public boolean check(T obj);
    }
}
